"""
模型加密/解密核心模块 - V4.0零信任架构
提供安全的模型文件加密和解密功能，防止模型知识产权泄露

安全特性:
- 内存安全的加密/解密
- 分块处理大文件
- 完整性验证
- 安全内存清理
"""

import os
import io
import hashlib
import base64
import structlog
from typing import Optional, Dict, Any, Tuple
from pathlib import Path
from cryptography.fernet import Fernet, InvalidToken
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import gc

logger = structlog.get_logger()


class ModelEncryptionError(Exception):
    """模型加密相关异常"""
    pass


class SecureModelLoader:
    """安全的模型加载器 - 支持内存解密和验证"""

    def __init__(self, encryption_key: str, chunk_size: int = 1024 * 1024):
        """
        初始化安全模型加载器

        Args:
            encryption_key: Base64编码的AES-256密钥
            chunk_size: 分块大小（字节）
        """
        self.chunk_size = chunk_size
        self._validate_and_setup_key(encryption_key)

    def _validate_and_setup_key(self, key_str: str) -> None:
        """验证并设置加密密钥"""
        try:
            self.key = key_str.encode('utf-8')
            self.fernet = Fernet(self.key)
            logger.info("加密密钥验证通过")
        except Exception as e:
            logger.error("加密密钥验证失败", error=str(e))
            raise ModelEncryptionError(f"无效的加密密钥: {e}")

    def load_encrypted_model(self, model_path: str) -> io.BytesIO:
        """
        加载并解密模型文件到内存

        Args:
            model_path: 加密的模型文件路径

        Returns:
            包含解密模型数据的内存缓冲区

        Raises:
            ModelEncryptionError: 解密失败时抛出
        """
        model_file = Path(model_path)

        if not model_file.exists():
            raise ModelEncryptionError(f"模型文件不存在: {model_path}")

        if not model_file.suffix == '.enc':
            raise ModelEncryptionError(f"模型文件扩展名应为.enc: {model_path}")

        try:
            logger.info("开始安全加载加密模型", model_path=model_path)

            # 验证文件完整性
            if not self._verify_file_integrity(model_file):
                raise ModelEncryptionError("模型文件完整性验证失败")

            # 在内存中解密模型
            decrypted_data = self._decrypt_in_memory(model_file)

            logger.info("模型解密成功",
                       file_size=model_file.stat().st_size,
                       data_size=len(decrypted_data))

            return io.BytesIO(decrypted_data)

        except InvalidToken:
            logger.error("模型解密失败 - 无效的密钥或损坏的文件")
            raise ModelEncryptionError("模型解密失败：密钥无效或文件已损坏")
        except Exception as e:
            logger.error("模型加载过程出错", error=str(e))
            raise ModelEncryptionError(f"模型加载失败: {e}")

    def _verify_file_integrity(self, file_path: Path) -> bool:
        """验证文件完整性"""
        try:
            # 检查文件大小
            file_size = file_path.stat().st_size
            if file_size < 100:  # 最小文件大小检查
                logger.error("文件过小，可能已损坏", file_size=file_size)
                return False

            # 检查文件头（Fernet加密文件应该有特定格式）
            with open(file_path, 'rb') as f:
                header = f.read(100)  # 读取前100字节

            # Fernet加密数据以特定格式开始
            if len(header) < 5:
                logger.error("文件头不完整")
                return False

            # 验证加密数据格式
            try:
                # 尝试解密第一个小块验证格式
                self.fernet.decrypt(header[:100])
            except InvalidToken:
                logger.error("文件格式验证失败")
                return False
            except Exception:
                # 其他异常可能是数据不足，这是正常的
                pass

            return True

        except Exception as e:
            logger.error("文件完整性检查失败", error=str(e))
            return False

    def _decrypt_in_memory(self, file_path: Path) -> bytes:
        """
        在内存中解密文件（安全地处理大文件）

        Args:
            file_path: 加密文件路径

        Returns:
            解密后的数据
        """
        decrypted_chunks = []

        try:
            with open(file_path, 'rb') as f:
                chunk_count = 0
                total_decrypted = 0

                while True:
                    # 读取加密分块
                    encrypted_chunk = f.read(self.chunk_size)
                    if not encrypted_chunk:
                        break

                    # 解密分块
                    decrypted_chunk = self.fernet.decrypt(encrypted_chunk)
                    decrypted_chunks.append(decrypted_chunk)

                    chunk_count += 1
                    total_decrypted += len(decrypted_chunk)

                    # 清理中间变量
                    del encrypted_chunk

                    # 定期触发垃圾回收
                    if chunk_count % 50 == 0:
                        gc.collect()

                # 合并所有分块
                decrypted_data = b''.join(decrypted_chunks)

                logger.info("文件解密完成",
                           total_chunks=chunk_count,
                           total_size=total_decrypted)

                return decrypted_data

        except Exception as e:
            # 清理已解密的数据
            for chunk in decrypted_chunks:
                del chunk
            decrypted_chunks.clear()
            gc.collect()

            logger.error("内存解密过程失败", error=str(e))
            raise

    def secure_cleanup(self, data: bytes) -> None:
        """安全清理敏感数据"""
        try:
            # 覆写内存中的敏感数据
            if isinstance(data, bytearray):
                for i in range(len(data)):
                    data[i] = 0
            # 强制垃圾回收
            gc.collect()
            logger.debug("敏感数据已安全清理")
        except Exception as e:
            logger.warning("数据清理失败", error=str(e))


class ModelIntegrityValidator:
    """模型完整性验证器"""

    @staticmethod
    def calculate_hash(data: bytes) -> str:
        """计算数据的SHA-256哈希"""
        return hashlib.sha256(data).hexdigest()

    @staticmethod
    def validate_model_structure(data: bytes) -> Dict[str, Any]:
        """
        验证模型结构（基础检查）

        Args:
            data: 模型数据

        Returns:
            验证结果字典
        """
        validation_result = {
            'is_valid': True,
            'issues': [],
            'file_size': len(data),
            'estimated_parameters': None
        }

        try:
            # 基础结构检查
            if len(data) < 1024:  # 至少1KB
                validation_result['issues'].append("文件过小，可能不是有效的模型文件")
                validation_result['is_valid'] = False

            # 检查是否为有效的pickle数据
            import pickle
            try:
                # 尝试反序列化（不实际加载，只验证格式）
                pickle.loads(data)
            except Exception as e:
                if "invalid load key" in str(e):
                    validation_result['issues'].append("无效的模型文件格式")
                    validation_result['is_valid'] = False

            # 估算参数数量（简化版本）
            estimated_params = len(data) // 4  # 粗略估算
            validation_result['estimated_parameters'] = estimated_params

            if validation_result['issues']:
                logger.warning("模型结构验证发现问题", issues=validation_result['issues'])
            else:
                logger.info("模型结构验证通过")

        except Exception as e:
            validation_result['is_valid'] = False
            validation_result['issues'].append(f"验证过程异常: {e}")
            logger.error("模型结构验证失败", error=str(e))

        return validation_result


class EncryptionKeyManager:
    """加密密钥管理器"""

    @staticmethod
    def derive_key_from_password(password: str, salt: Optional[bytes] = None) -> Tuple[bytes, bytes]:
        """
        从密码派生密钥

        Args:
            password: 用户密码
            salt: 盐值，如果为None则随机生成

        Returns:
            (derived_key, salt)
        """
        if salt is None:
            salt = os.urandom(16)

        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )

        key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
        return key, salt

    @staticmethod
    def validate_key_strength(key: str) -> bool:
        """验证密钥强度"""
        try:
            key_bytes = key.encode('utf-8')
            fernet = Fernet(key_bytes)
            return len(key_bytes) == 44  # Fernet密钥固定长度
        except Exception:
            return False


# 全局实例（延迟初始化）
_secure_loader: Optional[SecureModelLoader] = None


def get_secure_loader(encryption_key: str) -> SecureModelLoader:
    """获取安全模型加载器实例"""
    global _secure_loader
    if _secure_loader is None:
        _secure_loader = SecureModelLoader(encryption_key)
    return _secure_loader


def load_model_securely(model_path: str, encryption_key: str) -> io.BytesIO:
    """
    安全加载加密模型的便捷函数

    Args:
        model_path: 加密模型文件路径
        encryption_key: 解密密钥

    Returns:
        包含解密模型数据的内存缓冲区
    """
    loader = get_secure_loader(encryption_key)
    return loader.load_encrypted_model(model_path)


# 导出主要类和函数
__all__ = [
    'ModelEncryptionError',
    'SecureModelLoader',
    'ModelIntegrityValidator',
    'EncryptionKeyManager',
    'get_secure_loader',
    'load_model_securely'
]